# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import json
from dataclasses import dataclass
from collections import OrderedDict
from typing import Dict, Optional
from pathlib import Path

# "Qwen2-0.5B":
SUPPORTED_MODELS = OrderedDict()


@dataclass
class ModelMetadata:
    template: str
    path: Dict[str, str]


def register_model_group(
    models: Dict[str, Dict[str, str]],
    template: Optional[str] = None,
) -> None:
    """
    Register a model group with the given models

    Args:
        model_group (str): Name of the model group.
        models (Dict[str, Dict[str, str]]): A dict of models, where the key is the model name and the value is the model path.
        template (Optional[str], optional): The prompt template of the model group. Defaults to None.
    """

    for name, path in models.items():
        SUPPORTED_MODELS[name] = ModelMetadata(template, path)


def register_builtin_models():
    """
    Registers all builtin models which are predefined in 'configs/models.json'.
    """
    models_path = Path(__file__).resolve().parent.parent / "configs/models.json"
    with open(models_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    for _, model_item in config.items():
        register_model_group(
            models=model_item["models"],
            template=model_item.get("template"),
        )


register_builtin_models()
